#define pexit(s) { fprintf( stderr, "%s\n", s ); exit(1); }
#include <vector>
#include <string>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <dirent.h>
#include "jerasure.h"
#include "reed_sol.h"
#include "galois.h"
using namespace std;

int oct_to_int(char *oct);
vector<string> read_directory( string dir );
unsigned int compute_checksum( unsigned char* buffer, unsigned int blocksize );

int main( int argc, char** argv ){

	FILE* f_new_tar;
	FILE* f_pc;
	string dir;
	vector<string> files;
	vector<unsigned int> tarred_sizes;
	unsigned int fs_orig, fs_tarred; //size of a file going into the archive, original and tar'd
	string tarcommand;
	int failed_data, failed_coding;
	int i, j, l, s;
	unsigned char *buffer, *buffer2;
	unsigned char tar_buf[512];
	unsigned int checkbuffer;
	int k, m, w, blocksize, cursize;
	char coding_method[100];
	JER_Slices *slices;
	JER_Gen_T *gen;
	JER_Matrix *decode_matrix;
	uint64_t val;
	int* dmids;
	vector <int> is_dmid;

	printf("\n");

	if( argc != 3 ){
		pexit( "usage: ./personal_retrieval dir pc\n" );
	}
	dir = string( argv[1] );

	/* Open coding file and read in n, k, w, blocksize. Determine what method 
	   was used in encoding. Create slices object to include appropriate 
	   generator object. */
	f_pc = fopen( argv[2], "r" );	
	fscanf( f_pc, "%d %d %d %s %d\n", &k, &m, &w, coding_method, &blocksize );

	printf("Retrieved coding information:\n");
	printf("%5sk=%d\n","",k);
	printf("%5sm=%d\n","",m);
	printf("%5sw=%d\n","",w);
	printf("%5smethod=%s\n","",coding_method);
	printf("%5sblocksize=%d\n\n","",blocksize);

	/* Create files vector. */
	files = read_directory( dir );
	files.insert( files.begin(), string(dir) );
	printf("Tar'ing the input folder\n\n");
	tarcommand = "tar -cf temp.tar " + dir;
	system( tarcommand.c_str() );

	f_new_tar = fopen( "temp.tar", "r+" );
	if(f_new_tar == NULL){
		fprintf(stderr,"Error: could not open tar file\n");
		exit(1);
	}

	printf("Listing tar file's contents (note: All files are zero-padded to a multiple of 512 bytes. Each file also contains 512 extra bytes of meta-data.)\n");
	while(!feof(f_new_tar)){
		fs_tarred = fread( tar_buf, 1, 512, f_new_tar );
		if( !feof( f_new_tar ) && tar_buf[0] != 0 ){
			//read the original size of the file
			fs_orig = oct_to_int((char *)tar_buf+124);

			if( fs_orig != 0 )
				fs_tarred += ( ( fs_orig - 1 ) / 512 + 1 ) * 512;
			
			tarred_sizes.push_back(fs_tarred );
			printf( "%5s%-20s Original size: %-10d Size in .tar: %-10d\n","",tar_buf,fs_orig, fs_tarred);
			fseek( f_new_tar, fs_tarred-512, SEEK_CUR );
		}
	}
	printf("\n");
	fseek( f_new_tar, 0, SEEK_SET );

	printf("Creating the generator matrix\n\n");
	gen = RS_Classic_Generator( k, m, w );
	if(gen == NULL){
		fprintf(stderr,"Error creating generator\n");
		exit(1);
	}

	slices = new JER_Slices( k + m, k, blocksize, 1, gen );

	/* Read coding blocks. */
	buffer2 = (unsigned char*) malloc( blocksize * sizeof(char) );
	buffer = (unsigned char*) malloc( blocksize * m * sizeof(char) );
	fread( buffer, m, blocksize, f_pc );

	/* Compare checksums of the original data against stored checksums.
	   If they don't match set that block to failed. */
	failed_data = 0;
	printf("Comparing data checksums\n");
	for( i = 0, l = 0; i < files.size(); i++ ){

		cursize = tarred_sizes[i];
		while( cursize > 0 ){
			bzero( buffer2, blocksize );
			cursize -= fread( buffer2, 1, 
					(blocksize < cursize) ? blocksize : cursize, f_new_tar );
			fscanf( f_pc, "%u ", &checkbuffer );
			if( compute_checksum( buffer2, blocksize ) != checkbuffer ){
				slices->States[l] = 1;
				failed_data++;
			}else{
				slices->States[l] = 0;
			}
			l++;
		}
	}

	/* Compare checksums of coding blocks against stored checksums.
	   If they don't match, set that block to failed. */
	failed_coding = 0;
	printf("Comparing coding checksums\n");
	for( i = 0, l = k; i < m; i++ ){
		//read the checksum of the stored coding block
		fscanf( f_pc, "%u ", &checkbuffer );
		if( compute_checksum( buffer+i*blocksize, blocksize ) != checkbuffer ){
			slices->States[l] = 1;
			failed_coding++;
		}else{
			slices->States[l] = 0;
		}
		l++;
	}

	printf("\n");
	printf("Number of corrupt data blocks: %d\n",failed_data);
	printf("Number of corrupt coding blocks: %d\n",failed_coding);
	printf("\n");

	//need k survivors
	if(failed_data+failed_coding > m){
		fprintf(stderr,"Error: Cannot recover from %d failures with m = %d\n",failed_data+failed_coding,m);
		exit(1);
	}else if(failed_data+failed_coding == 0){
		printf("No corrupt blocks to decode\n");
		return 0;
	}

	//Make Decoding Matrix, fills in dmids
	dmids = (int*) malloc( sizeof(int) * k );
	is_dmid.resize(k+m,0);
	decode_matrix = slices->Make_Decoding_Matrix( dmids );
	if(decode_matrix == NULL){
		fprintf(stderr,"Error: could not create decoding matrix\n");
		exit(1);
	}
	//setup is_dmid so we can quickly check if a given block is a dmid
	for(i=0;i<k;i++){
		is_dmid[dmids[i]] = 1;
	}

	// Assign allocated memory to each failed drive. 
	for( j = 0, l = 0; j < k+m; j++ ){
		
		if( slices->States[j] == 1 ){
			//if a block failed, give it space
			slices->Ptrs[j] = buffer + l * blocksize;
			l++;
		}
		else{
			slices->Ptrs[j] = NULL;
		}
	}

	//zero out all the failed drives so that we can XOR into it 
	bzero( buffer, m * blocksize );

	/* 
		 Go through each file and process each block within the file.
	   If the block has failed, ignore it.
		 If the block has not failed and is in dm_ids, use it to decode
		 all of the failed data drives.
	*/

	fseek( f_new_tar, 0, SEEK_SET );

	//for each file
	for( i = 0, s = 0, l = 0; i < files.size(); i++ ){
		//note: l is the block id
		
		cursize = tarred_sizes[i];
		
		//get each block
		while( cursize > 0 ){
			bzero( buffer2, blocksize );
			cursize -= fread( buffer2, 1, (blocksize < cursize) ? blocksize : cursize, f_new_tar );

			//if the block is a survivor from our decoding matrix
			if(is_dmid[l] == 1){

				//use it to decode all the blocks that failed
				for( j = 0; j < k; j++ ){
					
					if( slices->States[j] == 1 ){
						//this block failed

						val = decode_matrix->Get( j, s );

						/*
						printf("block %d is:\n",l);
						for(int z = 0; z<100;z++){
							printf("%c",buffer2[z]);
						}
						printf("\n\n");
						printf("7th block is:\n");
						for(int z = 0; z<100;z++){
							printf("%c",slices->Ptrs[7][z]);
						}
						cout << "7th block, 43rd byte = " << slices->Ptrs[7][43] << endl;
						printf("\n\n");
						printf("decoding matrix val is: %llu\n",val);
						*/

						if( val == 1 ){
							galois_region_xor( slices->Ptrs[j], buffer2, slices->Ptrs[j], blocksize );
						}else{
							switch( w )
							{
								case 8:
									galois_w08_region_multiply( buffer2, val, 
											blocksize, slices->Ptrs[j], 1 );
									break;
								case 16:
									galois_w16_region_multiply( buffer2, val, 
											blocksize, slices->Ptrs[j], 1 );
									break;
								case 32:
									galois_w32_region_multiply( buffer2, val, 
											blocksize, slices->Ptrs[j], 1 );
									break;
							}
						}
					}
				}
				s++;
			}
			l++;
		}
	}
	/* Process must be applied to coding blocks as well. Read from beginning
	   of the coding file. Ignore the first line. Read in each block and
	   multiply by the appropriate matrix element and xor. */
	fseek( f_pc, 0, SEEK_SET );
	fscanf( f_pc, "%d %d %d %s %d\n", &k, &m, &w, coding_method, &blocksize );
	//for each coding block
	for( l = k; l < k + m; l++ )
	{
		fread( buffer2, 1, blocksize, f_pc );
		//if this coding block did not fail
		if( is_dmid[l] == 1)
		{
				//use it to decode all other data blocks that failed
				for( j = 0; j < k; j++ )
				{
					if( slices->States[j] == 1 )
					{
						val = decode_matrix->Get( j, s );

						if( val == 1 )
							galois_region_xor( slices->Ptrs[j], buffer2,
											   slices->Ptrs[j], blocksize );
						else
						{
							switch( w )
							{
								case 8:
									galois_w08_region_multiply( buffer2, val, 
											blocksize, slices->Ptrs[j], 1 );
									break;
								case 16:
									galois_w16_region_multiply( buffer2, val, 
											blocksize, slices->Ptrs[j], 1 );
									break;
								case 32:
									galois_w32_region_multiply( buffer2, val, 
											blocksize, slices->Ptrs[j], 1 );
									break;
							}
						}
					}
				}
				s++;
		}
	}

	//write the decoded file blocks to the tar file
	fseek( f_new_tar, 0, SEEK_SET );
	for( i = 0, l=0 ; i < files.size(); i++ )
	{
		cursize = tarred_sizes[i];
		while( cursize > 0 )
		{
			bzero( buffer2, blocksize );
			if( slices->Ptrs[l] != NULL ){
				//we decoded this block
				cursize -= fwrite( slices->Ptrs[l], 1, 
					(blocksize < cursize) ? blocksize : cursize, f_new_tar );
			}else{
				cursize -= fread( buffer2, 1, 
					(blocksize < cursize) ? blocksize : cursize, f_new_tar );
			}
			l++;
		}
	}
	fclose(f_new_tar);

	//untar
	system( "tar xvf temp.tar" );
}

unsigned int compute_checksum( unsigned char* buffer, unsigned int blocksize )
{
	unsigned int i, retval;

	retval = 0;
	for( i = 0; i < blocksize; i++ )
	{
		retval += buffer[i];
	}

	return retval;
}

vector<string> read_directory( string dir )
{
	vector<string> files;
	vector<string> tfiles;
	string s1, s2;
	DIR* dp;
	struct stat st_obj;
	struct dirent *dirp;
	int i;

	if ( ( dp = opendir( dir.c_str() ) ) == NULL )
	{
		fprintf( stderr, "Error: could not open directory %s\n", dir.c_str() );
		exit(1);
	}

	while( ( dirp = readdir( dp ) ) != NULL )
	{
		if( strcmp( dirp->d_name, "." ) && strcmp( dirp->d_name, ".." ) )
		{
			s1 = string( dirp->d_name );
			s2 = dir + "/" + s1;
			if( stat( s2.c_str(), &st_obj ) )
				pexit( "Error: in calling stat" );
			
			files.push_back( s2 );
			if( S_ISDIR( st_obj.st_mode ) )
			{
				tfiles = read_directory( s2 );
				for( i = 0; i < tfiles.size(); i++ )
					files.push_back( tfiles[i] );
			}
		}
	}
	closedir( dp );
	return files;
}

int oct_to_int(char *oct){
	int i;

	sscanf(oct, "%o", &i);

	return i;
}
